package com.yl.flink.sink;

import com.yl.constant.RedisCmd;
import com.yl.flink.redis.RedisMapper;
import com.yl.flink.redis.RedisCmdContainerBuilder;
import com.yl.flink.redis.RedisCmdDesc;
import com.yl.flink.redis.RedisContainer;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.streaming.api.functions.sink.RichSinkFunction;
import org.apache.flink.streaming.connectors.redis.RedisSink;
import org.apache.flink.streaming.connectors.redis.common.config.FlinkJedisConfigBase;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.IOException;
import java.util.Objects;

/**
 * @author wlf
 * @since 2022/8/25
 */
public class CustomRedisSink<IN> extends RichSinkFunction<IN> {

    private static final long serialVersionUID = 1L;
    private static final Logger LOG = LoggerFactory.getLogger(RedisSink.class);
    private final RedisCmd redisCommand;
    private final RedisMapper<IN> redisSinkMapper;
    private final FlinkJedisConfigBase flinkJedisConfigBase;
    private RedisContainer redisCommandsContainer;

    public CustomRedisSink(FlinkJedisConfigBase flinkJedisConfigBase, RedisMapper<IN> redisSinkMapper) {
        Objects.requireNonNull(flinkJedisConfigBase, "Redis connection pool config should not be null");
        Objects.requireNonNull(redisSinkMapper, "Redis Mapper can not be null");
        Objects.requireNonNull(redisSinkMapper.getCmdDesc(), "Redis Mapper data type description can not be null");
        this.flinkJedisConfigBase = flinkJedisConfigBase;
        this.redisSinkMapper = redisSinkMapper;
        RedisCmdDesc redisCommandDescription = redisSinkMapper.getCmdDesc();
        this.redisCommand = redisCommandDescription.getCommand();
    }

    @Override
    public void invoke(IN input, Context context) throws Exception {
        String key = this.redisSinkMapper.getKeyFromData(input);
        String value = this.redisSinkMapper.getValueFromData(input);
        int ttl = this.redisSinkMapper.getTtlFromData(input);
        switch (this.redisCommand) {
            case RPUSH:
                this.redisCommandsContainer.rpush(key, value);
                break;
            case LPUSH:
                this.redisCommandsContainer.lpush(key, value);
                break;
            case SET:
                this.redisCommandsContainer.set(key, value);
                break;
            case STTL:
                this.redisCommandsContainer.sttl(key, value, ttl);
                break;
            default:
                throw new IllegalArgumentException("Cannot process such data type: " + this.redisCommand);
        }
    }

    @Override
    public void open(Configuration parameters) throws Exception {
        try {
            this.redisCommandsContainer = RedisCmdContainerBuilder.build(this.flinkJedisConfigBase);
            this.redisCommandsContainer.open();
        } catch (Exception var3) {
            LOG.error("Redis has not been properly initialized: ", var3);
            throw var3;
        }
    }

    @Override
    public void close() throws IOException {
        if (this.redisCommandsContainer != null) {
            this.redisCommandsContainer.close();
        }
    }

}
